package class03

import "fmt"

/*
	请实现如下结构：
	NewTopRecord(k int) 构造时指定K的大小，构造后固定不变了
	TopRecord {
		add(str string)	 向该结构中加入一个字符串，可以重复加入
		top()            返回之前加入的所有字符串中，词频最大的K个
	}
	要求：
		add方法，复杂度O(log K)
		top方法，复杂度O(K)
	思路：
		建立大小为K的小根堆
		创建一个计数map
		创建一个记录位置的map
		计数map用于新加入的str的词频数
		位置map用于动态调整堆结构
*/
func NewTopRecord(K int) TopRecord {
	return TopRecord{
		freqMap:  make(map[string]*TopRecordNode),
		indexMap: make(map[*TopRecordNode]int),
		size:  K,
		index: 0,
		nodes: make([]*TopRecordNode, K),
	}
}

type TopRecordNode struct {
	Key string
	Num int
}

type TopRecord struct {
	freqMap  map[string]*TopRecordNode
	indexMap map[*TopRecordNode]int
	index    int
	size     int
	nodes    []*TopRecordNode
}

func (tp *TopRecord) add(s string) {
	// 频率计数
	if _, ok := tp.freqMap[s]; ok {
		tp.freqMap[s].Num++
	} else {
		tp.freqMap[s] = &TopRecordNode{
			Key: s,
			Num: 1,
		}
		// 首次出现必定无词频
		tp.indexMap[tp.freqMap[s]] = -1
	}

	// 查看频率是否出现在堆中
	if tp.indexMap[tp.freqMap[s]] == -1 {
		// 没有出现时，需要查看是否可以将顶点弹出，自身进入堆中
		if tp.index < tp.size { // 当堆没有放满的时候，放入堆中
			tp.Insert(tp.freqMap[s])
		} else {
			if tp.freqMap[s].Num > tp.peek().Num { // 当前词频大于堆顶
				tp.Poll()
				tp.Insert(tp.freqMap[s])
			}
		}
	} else {
		// 在堆中的话，计数改变，需要调整位置，看看是否需要下落
		tp.heapify(tp.indexMap[tp.freqMap[s]])
	}
}

func (tp *TopRecord) Poll() *TopRecordNode {
	if tp.index == 0 {
		panic(any("heap is empty!"))
	}
	tp.swap(0, tp.index-1)
	tp.index--
	tp.heapify(0)
	tp.indexMap[tp.nodes[tp.index]] = -1
	return tp.nodes[tp.index]
}

func (tp *TopRecord) heapify(index int) {
	parent := index
	small := 2*parent + 1
	for small < tp.index {
		if small+1 < tp.index && tp.nodes[small].Num > tp.nodes[small+1].Num {
			small = small+1
		}
		if tp.nodes[parent].Num >  tp.nodes[small].Num {
			tp.swap(parent, small)
		} else {
			break
		}
		parent = small
		small = 2*parent + 1
	}
}

func (tp *TopRecord) Insert(node *TopRecordNode) {
	tp.nodes[tp.index] = node
	tp.indexMap[node] = tp.index
	cur := tp.index
	parent := (cur-1) / 2
	for cur > 0 {
		if tp.nodes[cur].Num < tp.nodes[parent].Num {
			tp.swap(cur, parent)
		} else {
			break
		}
		cur = parent
		parent = (cur-1) / 2
	}
	tp.index++
}

func (tp *TopRecord) swap(index1, index2 int) {
	tp.indexMap[tp.nodes[index1]], tp.indexMap[tp.nodes[index2]] = 	tp.indexMap[tp.nodes[index2]], tp.indexMap[tp.nodes[index1]]
	tp.nodes[index1], tp.nodes[index2] = tp.nodes[index2], tp.nodes[index1]
}

func (tp *TopRecord) peek() *TopRecordNode {
	if tp.index == 0 {
		panic(any("heap is empty!"))
	}
	return tp.nodes[0]
}

func (tp *TopRecord) top() {
	// 堆不空的话，遍历堆中的数据，并打印。
	index := 0
	for index < tp.index {
		fmt.Println(tp.nodes[index].Key)
		index++
	}
}
